#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@NAME          : id_generator.py
@TIME          : 2025/03/02 13:33:11
@AUTHOR        : chenlip
@VERSION       : 0.0.1
@DESCRIPTION   : id生成器
'''
# --------------------------引用--------------------------
import time, random
import threading
# --------------------------引用--------------------------

class SnowflakeIdGenerator:
    def __init__(self, machine_id=0, datacenter_id=0):
        """
        初始化雪花ID生成器
        
        Args:
            machine_id (int): 机器ID (0-31)
            datacenter_id (int): 数据中心ID (0-31)
        """
        # 时间戳部分、数据中心部分、机器ID部分、序列号部分的位数
        self.timestamp_bits = 21
        self.datacenter_bits = 2
        self.machine_bits = 3
        self.sequence_bits = 6
        
        # 各部分最大值
        self.max_machine_id = -1 ^ (-1 << self.machine_bits)
        self.max_datacenter_id = -1 ^ (-1 << self.datacenter_bits)
        self.max_sequence = -1 ^ (-1 << self.sequence_bits)
        
        # 偏移量
        self.machine_id_shift = self.sequence_bits
        self.datacenter_id_shift = self.sequence_bits + self.machine_bits
        self.timestamp_shift = self.sequence_bits + self.machine_bits + self.datacenter_bits
        
        # 开始时间戳 (2025-01-01 00:00:00 UTC)
        self.twepoch = 1735689600000
        
        # 属性赋值
        if machine_id > self.max_machine_id or machine_id < 0:
            raise ValueError(f"机器ID超出范围(0-{self.max_machine_id})")
        if datacenter_id > self.max_datacenter_id or datacenter_id < 0:
            raise ValueError(f"数据中心ID超出范围(0-{self.max_datacenter_id})")
            
        self.machine_id = machine_id
        self.datacenter_id = datacenter_id
        self.sequence = 0
        self.last_timestamp = -1
        self.lock = threading.Lock()
        
    def _gen_timestamp(self):
        """获取当前时间戳"""
        return int(time.time() * 1000)
        
    def _til_next_millis(self, last_timestamp):
        """等待到下一毫秒"""
        timestamp = self._gen_timestamp()
        while timestamp <= last_timestamp:
            timestamp = self._gen_timestamp()
        return timestamp
        
    def generate(self):
        """生成雪花ID"""
        with self.lock:
            timestamp = self._gen_timestamp()
            
            # 时间回拨问题处理
            if timestamp < self.last_timestamp:
                # 实际生产环境可能需要更复杂的处理
                timestamp = self._til_next_millis(self.last_timestamp)
                
            # 同一毫秒内序列号递增
            if timestamp == self.last_timestamp:
                self.sequence = (self.sequence + 1) & self.max_sequence
                # 同一毫秒序列号用尽
                if self.sequence == 0:
                    timestamp = self._til_next_millis(self.last_timestamp)
            else:
                self.sequence = 0
                
            self.last_timestamp = timestamp
            
            # 组装雪花ID
            snowflake_id = (
                ((timestamp - self.twepoch) << self.timestamp_shift) |
                (self.datacenter_id << self.datacenter_id_shift) |
                (self.machine_id << self.machine_id_shift) |
                self.sequence
            )
            
            return snowflake_id
        

class CompactIdGenerator:
    def __init__(self):
        self.counter = 0
        self.lock = threading.Lock()
        self.start_time = int(time.time() * 1000)
        self.last_timestamp = -1
        
    def generate(self):
        """生成紧凑的ID"""
        with self.lock:
            timestamp = int(time.time() * 1000)
            
            # 确保递增
            if timestamp <= self.last_timestamp:
                timestamp = self.last_timestamp + 1
                
            self.last_timestamp = timestamp
            
            # 增加计数器
            self.counter = (self.counter + 1) % 1000
            
            # 组合时间差(相对于启动时间)和计数器
            time_diff = timestamp - self.start_time
            
            # 生成10位数字ID: 时间差(7位) + 计数器(3位)
            numeric_id = time_diff * 1000 + self.counter
            
            return numeric_id
        
class TimeRandomIdGenerator:
    def __init__(self):
        self.lock = threading.Lock()
        self.last_time = -1
        self.counter = random.randint(0, 999)
    
    def generate(self):
        """生成基于时间的随机ID"""
        with self.lock:
            # 获取当前秒级时间戳
            current_time = int(time.time())
            
            # 如果是同一秒，增加计数器
            if current_time == self.last_time:
                self.counter = (self.counter + 1) % 1000
            else:
                self.counter = random.randint(0, 999)
                self.last_time = current_time
            
            # 组合: 秒级时间戳 (10位) + 3位随机/计数器
            # 最终生成10-13位数字ID
            return current_time * 1000 + self.counter